from fastapi import APIRouter, HTTPException
from tasks.celery_app import celery_app
import redis
from config.celery_config import CeleryConfig

router = APIRouter()

@router.get("/workers", summary="获取Celery Worker列表")
def get_workers():
    """获取所有Celery Worker的状态"""
    try:
        inspect = celery_app.control.inspect()
        
        # 获取worker统计信息
        stats = inspect.stats() or {}
        # 获取活跃的worker
        active = inspect.active() or {}
        # 获取注册的任务
        registered = inspect.registered() or {}
        
        workers = []
        for worker_name in stats.keys():
            worker_stats = stats.get(worker_name, {})
            worker_info = {
                "name": worker_name,
                "status": "online",
                "active_tasks": len(active.get(worker_name, [])),
                "registered_tasks": len(registered.get(worker_name, [])),
                "total_processed": worker_stats.get('total', {}).get('total', 0),
                "pool_size": worker_stats.get('pool', {}).get('max-concurrency', 0),
                "load": worker_stats.get('load', [0, 0, 0])
            }
            workers.append(worker_info)
        
        return {
            "total_workers": len(workers),
            "workers": workers
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取Worker信息失败: {str(e)}")

@router.get("/queues", summary="获取Celery队列状态")
def get_queues():
    """获取所有Celery队列的状态"""
    try:
        inspect = celery_app.control.inspect()
        
        # 获取活跃队列
        active_queues = inspect.active_queues() or {}
        # 获取队列统计
        stats = inspect.stats() or {}
        
        queues = []
        queue_stats = {}
        
        # 从配置中获取队列列表
        configured_queues = ['file_processing', 'vector_processing', 'batch_processing', 'default']
        
        for queue_name in configured_queues:
            queue_info = {
                "name": queue_name,
                "workers": [],
                "total_messages": 0,
                "active_tasks": 0
            }
            
            # 统计每个队列的worker和任务
            for worker_name, worker_queues in active_queues.items():
                for queue in worker_queues:
                    if queue['name'] == queue_name:
                        queue_info["workers"].append(worker_name)
                        queue_info["total_messages"] += queue.get('messages', 0)
            
            # 统计活跃任务
            active_tasks = inspect.active() or {}
            for worker_name, tasks in active_tasks.items():
                for task in tasks:
                    if task.get('delivery_info', {}).get('routing_key') == queue_name:
                        queue_info["active_tasks"] += 1
            
            queues.append(queue_info)
        
        return {
            "total_queues": len(queues),
            "queues": queues
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取队列信息失败: {str(e)}")

@router.get("/active-tasks", summary="获取活跃任务")
def get_active_tasks():
    """获取所有正在执行的任务"""
    try:
        inspect = celery_app.control.inspect()
        
        # 获取活跃任务
        active = inspect.active() or {}
        # 获取已注册的任务
        registered = inspect.registered() or {}
        
        all_tasks = []
        total_active = 0
        
        for worker_name, tasks in active.items():
            for task in tasks:
                task_info = {
                    "worker": worker_name,
                    "id": task.get('id', ''),
                    "name": task.get('name', ''),
                    "args": task.get('args', []),
                    "kwargs": task.get('kwargs', {}),
                    "state": task.get('state', 'STARTED'),
                    "time_start": task.get('time_start', 0),
                    "acknowledged": task.get('acknowledged', False),
                    "delivery_info": task.get('delivery_info', {}),
                    "queue": task.get('delivery_info', {}).get('routing_key', 'default')
                }
                all_tasks.append(task_info)
                total_active += 1
        
        return {
            "total_active_tasks": total_active,
            "active_tasks": all_tasks
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取活跃任务失败: {str(e)}")

@router.get("/stats", summary="获取Celery统计信息")
def get_stats():
    """获取Celery的详细统计信息"""
    try:
        inspect = celery_app.control.inspect()
        
        # 获取各种统计信息
        stats = inspect.stats() or {}
        active = inspect.active() or {}
        registered = inspect.registered() or {}
        revoked = inspect.revoked() or {}
        scheduled = inspect.scheduled() or {}
        
        # 计算总体统计
        total_workers = len(stats)
        total_active_tasks = sum(len(tasks) for tasks in active.values())
        total_registered_tasks = sum(len(tasks) for tasks in registered.values())
        total_revoked_tasks = sum(len(tasks) for tasks in revoked.values())
        total_scheduled_tasks = sum(len(tasks) for tasks in scheduled.values())
        
        # 计算总处理任务数
        total_processed = 0
        for worker_stats in stats.values():
            total_processed += worker_stats.get('total', {}).get('total', 0)
        
        return {
            "overview": {
                "total_workers": total_workers,
                "total_active_tasks": total_active_tasks,
                "total_registered_tasks": total_registered_tasks,
                "total_revoked_tasks": total_revoked_tasks,
                "total_scheduled_tasks": total_scheduled_tasks,
                "total_processed_tasks": total_processed
            },
            "worker_stats": stats,
            "active_tasks": active,
            "registered_tasks": registered,
            "revoked_tasks": revoked,
            "scheduled_tasks": scheduled
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取统计信息失败: {str(e)}")

@router.get("/redis-info", summary="获取Redis连接信息")
def get_redis_info():
    """获取Redis连接和队列信息"""
    try:
        # 解析Redis URL
        redis_url = CeleryConfig.BROKER_URL.replace("redis://", "")
        if "/" in redis_url:
            host_port, db = redis_url.split("/")
        else:
            host_port, db = redis_url, "0"
        
        if ":" in host_port:
            host, port = host_port.split(":")
        else:
            host, port = host_port, "6379"
        
        # 连接Redis
        r = redis.Redis(host=host, port=int(port), db=int(db), socket_connect_timeout=5)
        
        # 获取Redis信息
        info = r.info()
        
        # 获取队列长度
        queue_lengths = {}
        for queue_name in ['file_processing', 'vector_processing', 'batch_processing', 'default']:
            queue_lengths[queue_name] = r.llen(queue_name)
        
        return {
            "redis_connection": {
                "host": host,
                "port": int(port),
                "db": int(db),
                "connected": r.ping()
            },
            "redis_info": {
                "used_memory": info.get('used_memory', 0),
                "used_memory_human": info.get('used_memory_human', ''),
                "connected_clients": info.get('connected_clients', 0),
                "total_commands_processed": info.get('total_commands_processed', 0)
            },
            "queue_lengths": queue_lengths
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取Redis信息失败: {str(e)}")

@router.get("/task/{task_id}", summary="获取特定任务信息")
def get_task_info(task_id: str):
    """获取特定任务的详细信息"""
    try:
        # 获取任务结果
        result = celery_app.AsyncResult(task_id)
        
        task_info = {
            "id": task_id,
            "status": result.status,
            "ready": result.ready(),
            "successful": result.successful(),
            "failed": result.failed(),
            "info": result.info if result.info else None,
            "traceback": result.traceback if result.traceback else None
        }
        
        return task_info
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取任务信息失败: {str(e)}") 